from langchain_chroma import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain.schema import Document
import os
from typing import List, Optional


class VectorStoreManager:
    """向量数据库管理器"""

    def __init__(self):
        """初始化向量数据库管理器"""
        # 不指定任何自定义URL参数，使用默认值
        self.embedding_model = OllamaEmbeddings(model="deepseek-r1:14b")
        self.active_db = None

    def init_chroma(self, persist_directory: str):
        """初始化Chroma向量数据库

        Args:
            persist_directory: 向量数据库持久化目录

        Returns:
            初始化好的Chroma实例
        """
        # 确保目录存在
        os.makedirs(persist_directory, exist_ok=True)

        self.active_db = Chroma(
            persist_directory=persist_directory,
            embedding_function=self.embedding_model
        )
        return self.active_db

    def add_documents(self, documents: List[Document]):
        """添加文档到向量数据库

        Args:
            documents: 要添加的文档列表

        Raises:
            ValueError: 如果向量数据库尚未初始化
        """
        if not self.active_db:
            raise ValueError("请先初始化向量数据库")

        self.active_db.add_documents(documents)

    def get_retriever(self, search_kwargs: Optional[dict] = None):
        """获取检索器

        Args:
            search_kwargs: 检索参数，如k值等

        Returns:
            向量数据库检索器

        Raises:
            ValueError: 如果向量数据库尚未初始化
        """
        if not self.active_db:
            raise ValueError("请先初始化向量数据库")

        if search_kwargs is None:
            search_kwargs = {"k": 4}

        return self.active_db.as_retriever(search_kwargs=search_kwargs)
